package org.wlldTest.myUnet;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import org.dromara.easyai.config.UNetConfig;
import org.dromara.easyai.entity.ThreeChannelMatrix;
import org.dromara.easyai.matrixTools.Matrix;
import org.dromara.easyai.matrixTools.MatrixOperation;
import org.dromara.easyai.tools.ImageTools;
import org.dromara.easyai.tools.Picture;
import org.dromara.easyai.unet.UNetInput;
import org.dromara.easyai.unet.UNetManager;
import org.dromara.easyai.unet.UNetModel;
import org.dromara.easyai.yolo.YoloModel;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStreamWriter;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;

/**
 * @author lidapeng
 * @time 2025/3/3 20:12
 * @des unet测试类
 */
public class Test {
    public static UNetConfig unetConfig;
    public static UNetManager uNetManager;

    public static void main(String[] args) throws Exception {
        init();
        study();
        //test();
    }

    public static ThreeChannelMatrix getMyPicture(int x, int y, float value) throws Exception {
        Matrix matrixR = new Matrix(x, y);
        Matrix matrixG = new Matrix(x, y);
        Matrix matrixB = new Matrix(x, y);
        ThreeChannelMatrix threeChannelMatrix = new ThreeChannelMatrix();
        threeChannelMatrix.setX(x);
        threeChannelMatrix.setY(y);
        threeChannelMatrix.setMatrixR(matrixR);
        threeChannelMatrix.setMatrixG(matrixG);
        threeChannelMatrix.setMatrixB(matrixB);
        for (int i = 0; i < x; i++) {
            for (int j = 0; j < y; j++) {
                matrixR.setNub(i, j, value);
                matrixG.setNub(i, j, value);
                matrixB.setNub(i, j, value);
            }
        }
        return threeChannelMatrix;
    }

    public static void study2() throws Exception {
        ImageBack imageBack = new ImageBack();
        UNetInput uNetInput = uNetManager.getInput();
        ThreeChannelMatrix t = getMyPicture(50, 50, 0.9f);
        ThreeChannelMatrix e = getMyPicture(50, 50, 0.5f);
        for (int i = 0; i < 1560; i++) {
            System.out.println("i========" + i);
            uNetInput.postMessage(1, null, t, e, true);
        }
        uNetInput.postMessage(1, imageBack, t, null, false);
        ImageTools.writeImage(imageBack.getMatrix(), "/Users/lidapeng/job/testPic/t3.png");
    }

    public static void test() throws Exception {
        UNetInput uNetInput = uNetManager.getInput();
        ImageBack imageBack = new ImageBack();
        String test = "/Users/lidapeng/job/testPic/a.png";
        ThreeChannelMatrix tm = Picture.getThreeMatrix(test, false);
        uNetInput.postMessage(2, imageBack, tm, null, false);
        ImageTools.writeImage(imageBack.getMatrix(), "/Users/lidapeng/job/testPic/at.png");
    }

    public static void init() throws Exception {
        unetConfig = new UNetConfig();
        unetConfig.setXSize(533);//533
        unetConfig.setYSize(400);//400
        unetConfig.setCutting(false);
        unetConfig.setAuto(false);
        unetConfig.setGMaxTh(10000f);
//        unetConfig.setStudyRate(0.000001f);
//        unetConfig.setOneStudyRate(0.000001f);
        uNetManager = new UNetManager(unetConfig);
        uNetManager.insertModel(readCreatorModel(2));
    }

    public static void study() throws Exception {
        UNetInput uNetInput = uNetManager.getInput();
        ImageBack imageBack = new ImageBack();
        String test = "/Users/lidapeng/job/testPic/a.png";
        for (int k = 0; k < 1; k++) {
            for (int i = 1; i < 401; i++) {
                System.out.println("i===" + i + ",k===" + k);
                String url = "/Users/lidapeng/job/testPic/test/test" + i + ".png";
                String urlE = "/Users/lidapeng/job/testPic/test2/test" + i + ".png";
                ThreeChannelMatrix t = Picture.getThreeMatrix(url, false);
                ThreeChannelMatrix e = Picture.getThreeMatrix(urlE, false);
                uNetInput.postMessage(1, null, t, e, true);
            }
        }
        UNetModel uNetModel = uNetManager.getModel();
        writeModel(JSON.toJSONString(uNetModel), "/Users/lidapeng/job/testPic/cut2.json");
        ThreeChannelMatrix tm = Picture.getThreeMatrix(test, false);
        uNetInput.postMessage(2, imageBack, tm, null, false);
        ImageTools.writeImage(imageBack.getMatrix(), "/Users/lidapeng/job/testPic/at2.jpg");
    }

    private static void writeModel(String model, String url) throws IOException {//写出模型与 激活参数
        OutputStreamWriter out = new OutputStreamWriter(Files.newOutputStream(Paths.get(url)), StandardCharsets.UTF_8);
        out.write(model);
        out.close();
    }

    private static UNetModel readCreatorModel(int id) {
        File file = new File("/Users/lidapeng/job/testPic/cut" + id + ".json");
        String a = readPaper(file);
        return JSONObject.parseObject(a, UNetModel.class);
    }

    private static String readPaper(File file) {
        InputStream read = null;
        String context = null;
        try {
            read = Files.newInputStream(file.toPath());
            byte[] bt = new byte[read.available()];
            read.read(bt);
            context = new String(bt, StandardCharsets.UTF_8);
            read.close();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (read != null) {
                try {
                    read.close(); //确保关闭
                } catch (IOException el) {
                }
            }
        }
        return context;
    }

}
